import time
import heapq
from typing import List, Tuple, Dict, Set, Any, Optional
import copy

from env.grid_env import GridMap
from planner.dqn_planner import DQNPlanner
from cbs.constraint import Constraint, EdgeConstraint, ConstraintSet


class CBSNode:
    """
    CBS搜索树的节点
    """
    def __init__(self, constraint_set: ConstraintSet = None, solution: Dict[int, List[Tuple[int, int]]] = None, cost: int = 0):
        """
        初始化CBS节点
        
        Args:
            constraint_set: 约束集合
            solution: 当前解决方案，格式为 {agent_id: path}
            cost: 当前解决方案的总代价
        """
        self.constraint_set = constraint_set if constraint_set else ConstraintSet()
        self.solution = solution if solution else {}
        self.cost = cost
        
    def __lt__(self, other):
        """比较函数，用于优先队列排序"""
        return self.cost < other.cost


class CBSSolver:
    """
    CBS (Conflict-Based Search) 求解器
    使用DQN作为低层路径规划器
    """
    def __init__(self, grid_map: GridMap, model_path: str, config: Dict[str, Any]):
        """
        初始化CBS求解器
        
        Args:
            grid_map: 网格地图
            model_path: DQN模型路径
            config: 配置参数
        """
        self.grid_map = grid_map
        self.model_path = model_path
        self.config = config
        
        # 创建DQN规划器
        self.planner = DQNPlanner(grid_map, model_path)
        
        # CBS参数
        self.max_iterations = config.get('max_iterations', 1000)
        self.timeout = config.get('timeout', 30)  # 超时时间(秒)
    
    def find_solution(self, tasks: Dict[int, Tuple[Tuple[int, int], Tuple[int, int]]]) -> Dict[int, List[Tuple[int, int]]]:
        """
        为多个机器人找到无冲突的路径
        
        Args:
            tasks: 任务字典 {agent_id: (start, goal)}
            
        Returns:
            Dict[int, List[Tuple[int, int]]]: 路径字典 {agent_id: path}
        """
        # 记录开始时间
        start_time = time.time()
        
        # 创建根节点
        root = CBSNode()
        
        print(f"CBS开始规划 {len(tasks)} 个机器人的路径")
        
        # 为每个机器人计算初始路径
        for agent_id, (start, goal) in tasks.items():
            print(f"为机器人 {agent_id} 规划从 {start} 到 {goal} 的初始路径")
            # 使用plan_multiple_paths获取多条路径并选择最短的一条
            path = self.planner.plan_multiple_paths(start, goal, num_paths=5)
            if not path:
                print(f"警告: 无法为机器人 {agent_id} 找到从 {start} 到 {goal} 的路径")
                # 如果找不到路径，返回空路径
                path = [start]
            root.solution[agent_id] = path
            print(f"机器人 {agent_id} 初始路径长度: {len(path)}")
        
        # 计算根节点的代价
        root.cost = self.calculate_solution_cost(root.solution)
        
        # 初始化开放列表
        open_list = [root]
        
        # CBS主循环
        iterations = 0
        while open_list and iterations < self.max_iterations:
            iterations += 1
            
            if iterations % 100 == 0:
                print(f"CBS迭代次数: {iterations}, 开放列表长度: {len(open_list)}")
            
            # 检查是否超时
            elapsed_time = time.time() - start_time
            if elapsed_time > self.timeout:
                print(f"CBS求解超时 ({self.timeout}秒), 已完成 {iterations} 次迭代")
                # 返回当前最优解
                if open_list:
                    best_node = min(open_list, key=lambda node: node.cost)
                    return best_node.solution
                return {}
            
            # 获取代价最小的节点
            node = heapq.heappop(open_list)
            
            # 检查是否存在冲突
            conflict = self.find_conflict(node.solution)
            if not conflict:
                # 如果没有冲突，返回解决方案
                print(f"CBS找到无冲突解决方案，迭代次数: {iterations}, 耗时: {time.time() - start_time:.2f}秒")
                return node.solution
            
            # 解包冲突
            agent_i, agent_j, pos_i, pos_j, timestep = conflict
            print(f"发现冲突: 机器人 {agent_i} 和 {agent_j} 在时间步 {timestep} 发生冲突")
            
            # 创建两个子节点
            # 子节点1: 约束agent_i不能在timestep时刻到达pos_i
            child1 = CBSNode(copy.deepcopy(node.constraint_set), copy.deepcopy(node.solution))
            constraint1 = Constraint(pos_i, timestep, agent_i)
            child1.constraint_set.add_constraint(constraint1)
            
            # 重新规划agent_i的路径
            start, goal = tasks[agent_i]
            constraints = child1.constraint_set.get_constraints(agent_i)
            print(f"重新规划机器人 {agent_i} 的路径, 约束数量: {len(constraints) if constraints else 0}")
            # 使用plan_multiple_paths获取多条路径并选择最短的一条
            new_path = self.planner.plan_multiple_paths(start, goal, constraints, num_paths=3)
            
            if new_path:
                child1.solution[agent_i] = new_path
                child1.cost = self.calculate_solution_cost(child1.solution)
                heapq.heappush(open_list, child1)
                print(f"子节点1: 机器人 {agent_i} 新路径长度: {len(new_path)}")
            else:
                print(f"子节点1: 无法为机器人 {agent_i} 找到满足约束的路径")
            
            # 子节点2: 约束agent_j不能在timestep时刻到达pos_j
            child2 = CBSNode(copy.deepcopy(node.constraint_set), copy.deepcopy(node.solution))
            constraint2 = Constraint(pos_j, timestep, agent_j)
            child2.constraint_set.add_constraint(constraint2)
            
            # 重新规划agent_j的路径
            start, goal = tasks[agent_j]
            constraints = child2.constraint_set.get_constraints(agent_j)
            print(f"重新规划机器人 {agent_j} 的路径, 约束数量: {len(constraints) if constraints else 0}")
            # 使用plan_multiple_paths获取多条路径并选择最短的一条
            new_path = self.planner.plan_multiple_paths(start, goal, constraints, num_paths=3)
            
            if new_path:
                child2.solution[agent_j] = new_path
                child2.cost = self.calculate_solution_cost(child2.solution)
                heapq.heappush(open_list, child2)
                print(f"子节点2: 机器人 {agent_j} 新路径长度: {len(new_path)}")
            else:
                print(f"子节点2: 无法为机器人 {agent_j} 找到满足约束的路径")
        
        # 如果达到最大迭代次数或开放列表为空
        print(f"CBS无法找到完全无冲突的解决方案，迭代次数: {iterations}")
        
        # 返回当前最优解
        if open_list:
            best_node = min(open_list, key=lambda node: node.cost)
            return best_node.solution
        
        # 如果没有可行解，返回空字典
        return {}
    
    def find_conflict(self, solution: Dict[int, List[Tuple[int, int]]]) -> Optional[Tuple[int, int, Tuple[int, int], Tuple[int, int], int]]:
        """
        在解决方案中查找冲突
        
        Args:
            solution: 路径字典 {agent_id: path}
            
        Returns:
            Optional[Tuple[int, int, Tuple[int, int], Tuple[int, int], int]]: 
                冲突信息 (agent_i, agent_j, pos_i, pos_j, timestep)，如果没有冲突则返回None
        """
        # 获取所有机器人ID
        agent_ids = list(solution.keys())
        
        # 检查所有机器人对
        for i in range(len(agent_ids)):
            for j in range(i + 1, len(agent_ids)):
                agent_i = agent_ids[i]
                agent_j = agent_ids[j]
                
                path_i = solution[agent_i]
                path_j = solution[agent_j]
                
                # 获取两条路径的最大长度
                max_path_length = max(len(path_i), len(path_j))
                
                # 检查每个时间步
                for timestep in range(max_path_length):
                    # 获取两个机器人在当前时间步的位置
                    pos_i = path_i[min(timestep, len(path_i) - 1)]
                    pos_j = path_j[min(timestep, len(path_j) - 1)]
                    
                    # 检查顶点冲突（两个机器人在同一时间到达同一位置）
                    if pos_i == pos_j:
                        return (agent_i, agent_j, pos_i, pos_j, timestep)
                    
                    # 检查边冲突（两个机器人在同一时间交换位置）
                    if timestep > 0:
                        prev_pos_i = path_i[min(timestep - 1, len(path_i) - 1)]
                        prev_pos_j = path_j[min(timestep - 1, len(path_j) - 1)]
                        
                        if pos_i == prev_pos_j and pos_j == prev_pos_i:
                            return (agent_i, agent_j, pos_i, pos_j, timestep)
        
        # 没有找到冲突
        return None
    
    def calculate_solution_cost(self, solution: Dict[int, List[Tuple[int, int]]]) -> int:
        """
        计算解决方案的总代价（所有路径长度之和）
        
        Args:
            solution: 路径字典 {agent_id: path}
            
        Returns:
            int: 总代价
        """
        return sum(len(path) for path in solution.values())
    
    def get_makespan(self, solution: Dict[int, List[Tuple[int, int]]]) -> int:
        """
        计算解决方案的完工时间（最长路径的长度）
        
        Args:
            solution: 路径字典 {agent_id: path}
            
        Returns:
            int: 完工时间
        """
        return max(len(path) for path in solution.values()) if solution else 0 